import numpy as np
from sklearn import svm
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

from config import *



class DM_SVM:
    def __init__(self, x_train, x_test, y_train, y_test):
        self.x_train, self.x_test, self.y_train, self.y_test = (
            x_train,
            x_test,
            y_train,
            y_test,
        )
        self.model = [svm.SVR() for _ in out_feature_name_list]

    def train(self):
        for i, m in enumerate(self.model):
            m.fit(X=self.x_train, y=self.y_train[:, i])

    def predict(self):
        # for i, m in enumerate(self.model):
        #     s = m.score(self.x_test, self.y_test[:, i])
        #     print(s)
        pred = np.zeros(self.y_test.shape)
        acc = np.zeros(self.y_test.shape[1])
        for i, m in enumerate(self.model):
            pred[:, i] = m.predict(self.x_test)
            acc[i] = mean_squared_error(self.y_test[:, i], pred[:, i])

        print("SVM accuracy: ", acc)
        return acc, pred

    # def predict(self, inputs):
    #     pred = np.zeros(len(out_feature_name_list))
    #     inputs = np.array([inputs])
    #     for i, m in enumerate(self.model):
    #         pred[i] = m.predict(inputs)
    #     return pred
